# edited by levi    @触屏阈值选择器
# 阈值调用例子 print(thresholds["Red"])
from maix import touchscreen, app, display, image, camera
import time, os


# 全局变量
cam = camera.Camera(640, 480)
ts = touchscreen.TouchScreen()
disp = display.Display()
buttons = []
zoomed = False
thresholds = {
    "Red": [0, 0, 0, 0],    # 初始化Red阈值
    "Green": [0, 0, 0, 0],  # 初始化Green阈值
    "Blue": [0, 0, 0, 0]    # 初始化Blue阈值
}
selected_color = "Red"  # 当前选中的颜色
last_touch_time = 0  # 上次触摸时间
debounce_delay = 0.1  # 消抖延迟（秒）
img = cam.read()  # 初始化摄像头画面
show_threshold_start_time = 0  # 开始显示阈值的起始时间
show_threshold_duration = 0.5  # 显示阈值的持续时间（秒）
accumulated_statistics = {
    "Red": [0, 0, 0, 0, 0, 0],  # 初始化累积Red阈值统计信息
    "Green": [0, 0, 0, 0, 0, 0],  # 初始化累积Green阈值统计信息
    "Blue": [0, 0, 0, 0, 0, 0],  # 初始化累积Blue阈值统计信息
}
frame_count = 0  # 帧计数器

# 按钮类
def create_button(label, x, y, width=None, height=None, color=image.COLOR_WHITE):
    size = image.string_size(label)
    if width is None:
        width = 8 * 2 + size.width()
    if height is None:
        height = 12 * 2 + size.height()
    return {
        "label": label,
        "x": x,
        "y": y,
        "width": width,
        "height": height,
        "color": color,
        "size": size
    }

# 绘制按钮
def draw_button(button, img):
    img.draw_rect(button["x"], button["y"], button["width"], button["height"], button["color"], 2)
    img.draw_string(button["x"] + 8, button["y"] + 12, button["label"], button["color"])

# 检查按钮是否被按下
def is_button_pressed(button, x, y):
    return (button["x"] <= x <= button["x"] + button["width"] and
            button["y"] <= y <= button["y"] + button["height"])

# 加载阈值
def load_thresholds():
    file_path = "/root/thresholds.txt"
    if not os.path.exists(file_path):
        open(file_path, 'w').close()  # 如果文件不存在，创建一个新文件
    try:
        with open(file_path, "r") as file:
            lines = file.readlines()
            for line in lines:
                color, threshold = line.strip().split(": ")
                thresholds[color] = eval(threshold)
    except FileNotFoundError:
        print("Threshold file not found")

# 保存阈值
def save_thresholds():
    with open("/root/thresholds.txt", "w") as file:
        for color, threshold in thresholds.items():
            if threshold is not None:
                file.write(f"{color}: {threshold}\n")

# 微调阈值
def adjust_threshold(color, channel, amount):
    index = 2 * channel
    thresholds[color][index] += amount
    thresholds[color][index + 1] += amount
    if thresholds[color][index] < -255:
        thresholds[color][index] = -255
    if thresholds[color][index + 1] > 255:
        thresholds[color][index + 1] = 255

# 缩减阈值区间
def shrink_threshold(color, channel, amount):
    index = 2 * channel
    thresholds[color][index] += amount
    thresholds[color][index + 1] -= amount
    if thresholds[color][index] < -255:
        thresholds[color][index] = -255
    if thresholds[color][index + 1] > 255:
        thresholds[color][index + 1] = 255

# 扩大阈值区间
def expand_threshold(color, channel, amount):
    index = 2 * channel
    thresholds[color][index] -= amount
    thresholds[color][index + 1] += amount
    if thresholds[color][index] < -255:
        thresholds[color][index] = -255
    if thresholds[color][index + 1] > 255:
        thresholds[color][index + 1] = 255

# 动态阈值刷新函数
def refresh_thresholds(color, blobs):
    global frame_count, accumulated_statistics
    if not blobs:
        return
    # 计算色块的中心点
    centers = [(blob[0] + blob[2] / 2, blob[1] + blob[3] / 2) for blob in blobs]
    if not centers:
        return
    
    # 计算色块的平均中心点
    avg_x = sum(center[0] for center in centers) / len(centers)
    avg_y = sum(center[1] for center in centers) / len(centers)
    
    # 定义一个ROI区域
    roi_x, roi_y = int(avg_x - 10), int(avg_y - 10)
    roi_width, roi_height = 20, 20
    
    # 确保ROI在图像范围内
    if roi_x < 0: roi_x = 0
    if roi_y < 0: roi_y = 0
    if roi_x + roi_width > img.width(): roi_x = img.width() - roi_width
    if roi_y + roi_height > img.height(): roi_y = img.height() - roi_height
    
    # 获取ROI区域的统计信息
    statistics_data = img.get_statistics(roi=[roi_x, roi_y, roi_width, roi_height])
    
    # 每隔30帧重置累积统计信息和帧计数器
    if frame_count >= 5:
        accumulated_statistics[color] = [0, 0, 0, 0, 0, 0]
        frame_count = 0
        
    # 更新累积的颜色统计信息
    for i in range(6):
        new_value = (statistics_data.l_mode() - 15 if i % 2 == 0 else statistics_data.l_mode() + 15) if i < 2 else \
                    (statistics_data.a_mode() - 15 if i % 2 == 0 else statistics_data.a_mode() + 15) if i < 4 else \
                    (statistics_data.b_mode() - 15 if i % 2 == 0 else statistics_data.b_mode() + 15)
        accumulated_statistics[color][i] += new_value
    frame_count += 1

    # 计算平均阈值
    average_thresholds = [x / frame_count for x in accumulated_statistics[color]]
    
    # 平滑过渡
    for i in range(6):
        thresholds[color][i] += int(average_thresholds[i] - thresholds[color][i])

    print(f"Threshold set to LAB for {color}: {thresholds[color]}")
    save_thresholds()  # 保存阈值

# 运行主循环
def run():
    global img, selected_color, zoomed, last_touch_time, show_threshold_start_time
    while not app.need_exit():
        img = cam.read()  # 每次循环都读取新的摄像头画面

        for button in buttons:
            draw_button(button, img)  # 在每次读取摄像头画面后重新绘制按钮

        x, y, pressed = ts.read()
        current_time = time.time()
        if pressed and (current_time - last_touch_time > debounce_delay):
            last_touch_time = current_time
            for button in buttons:
                if is_button_pressed(button, x, y):
                    if button["label"] == "< Exit":
                        save_thresholds()  # 保存阈值
                        app.set_exit_flag(True)
                    elif button["label"] in ["Red", "Green", "Blue"]:
                        print(f"{button['label']} pressed!")
                        selected_color = button["label"]
                        zoomed = not zoomed  # 切换放大状态
                    elif button["label"] == "Show threshold":
                        load_thresholds()  # 加载阈值
                        show_threshold_start_time = current_time  # 记录开始显示阈值的时间
                    elif button["label"].startswith("L+") and selected_color:
                        adjust_threshold(selected_color, 0, 3)
                        print(f"Adjusted L threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"].startswith("L-") and selected_color:
                        adjust_threshold(selected_color, 0, -3)
                        print(f"Adjusted L threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"].startswith("A+") and selected_color:
                        adjust_threshold(selected_color, 1, 3)
                        print(f"Adjusted A threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"].startswith("A-") and selected_color:
                        adjust_threshold(selected_color, 1, -3)
                        print(f"Adjusted A threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"].startswith("B+") and selected_color:
                        adjust_threshold(selected_color, 2, 3)
                        print(f"Adjusted B threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"].startswith("B-") and selected_color:
                        adjust_threshold(selected_color, 2, -3)
                        print(f"Adjusted B threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"] == "S" and selected_color:
                        shrink_threshold(selected_color, 0, 1)
                        shrink_threshold(selected_color, 1, 1)
                        shrink_threshold(selected_color, 2, 1)
                        print(f"Shrank threshold for {selected_color}: {thresholds[selected_color]}")
                    elif button["label"] == "L" and selected_color:
                        expand_threshold(selected_color, 0, 1)
                        expand_threshold(selected_color, 1, 1)
                        expand_threshold(selected_color, 2, 1)
                        print(f"Expanded threshold for {selected_color}: {thresholds[selected_color]}")
                    break

        # 将屏幕坐标映射回图像坐标
        x, y = image.resize_map_pos_reverse(img.width(), img.height(), disp.width(), disp.height(), image.Fit.FIT_CONTAIN, x, y)
        if pressed and zoomed and selected_color and 30 < y < 410:
            # 如果处于放大模式，抓取触摸点的LAB阈值
            roi_x, roi_y = x - 10, y - 10  # 定义一个20x20的ROI区域
            if roi_x < 0: roi_x = 0
            if roi_y < 0: roi_y = 0
            if roi_x + 20 > img.width(): roi_x = img.width() - 20
            if roi_y + 20 > img.height(): roi_y = img.height() - 20

            statistics_data = img.get_statistics(roi=[roi_x, roi_y, 20, 20])
            thresholds[selected_color] = list(
                (
                    statistics_data.l_mode() - 15,  # L通道众数
                    statistics_data.l_mode() + 15,
                    statistics_data.a_mode() - 15,  # A通道众数
                    statistics_data.a_mode() + 15,
                    statistics_data.b_mode() - 15,  # B通道众数
                    statistics_data.b_mode() + 15
                )
            )
            print(f"Threshold set to LAB for {selected_color}: {thresholds[selected_color]}")
            save_thresholds()  # 保存阈值
                
        if selected_color and thresholds[selected_color] is not None and zoomed:
            # 如果设置了阈值且处于放大模式，对图像进行二值化处理
            img = img.binary([thresholds[selected_color]])
            for button in buttons:
                draw_button(button, img)  # 重新绘制按钮
        elif current_time - show_threshold_start_time < show_threshold_duration:
            # 如果处于显示阈值模式，并且时间未超过0.5秒，则继续显示当前颜色的二值化图像
            img = img.binary([thresholds[selected_color]])
            print(f"Threshold for {selected_color}: {thresholds}")
            for button in buttons:
                draw_button(button, img)  # 重新绘制按钮
        else:
            # 如果时间超过0.5秒
            img = cam.read()  # 读取新的摄像头画面
            for button in buttons:
                draw_button(button, img)  # 重新绘制按钮
        if thresholds["Red"] != [0, 0, 0, 0] and not zoomed:
            blobsr = img.find_blobs([thresholds["Red"]], area_threshold=2000, pixels_threshold=100, merge=True)
            if blobsr:
                refresh_thresholds("Red", blobsr)
            for blob in blobsr:
                img.draw_rect(blob[0], blob[1], blob[2], blob[3], image.COLOR_RED)
                cx = blob[0] + blob[2] / 2
                cy = blob[1] + blob[3] / 2
                img.draw_cross(int(cx), int(cy), image.Color.from_rgb(255, 0, 0), size=5, thickness=2)
        if thresholds["Green"] != [0, 0, 0, 0] and not zoomed:
            blobsg = img.find_blobs([thresholds["Green"]], area_threshold=2000, pixels_threshold=100, merge=True)
            if blobsg:
                refresh_thresholds("Green", blobsg)
            for blob in blobsg:
                img.draw_rect(blob[0], blob[1], blob[2], blob[3], image.COLOR_GREEN)
                cx = blob[0] + blob[2] / 2
                cy = blob[1] + blob[3] / 2
                img.draw_cross(int(cx), int(cy), image.Color.from_rgb(0, 255, 0), size=5, thickness=2)
        if thresholds["Blue"] != [0, 0, 0, 0] and not zoomed:
            blobsb = img.find_blobs([thresholds["Blue"]], area_threshold=2000, pixels_threshold=100, merge=True)
            if blobsb:
                refresh_thresholds("Blue", blobsb)
            for blob in blobsb:
                img.draw_rect(blob[0], blob[1], blob[2], blob[3], image.COLOR_BLUE)
                cx = blob[0] + blob[2] / 2
                cy = blob[1] + blob[3] / 2
                img.draw_cross(int(cx), int(cy), image.Color.from_rgb(0, 0, 255), size=5, thickness=2)
                
        disp.show(img, fit=image.Fit.FIT_CONTAIN)

if __name__ == "__main__":
    buttons.append(create_button("< Exit", 0, 0))
    buttons.append(create_button("Red", 100, 0, 65, color=image.COLOR_RED))
    buttons.append(create_button("Green", 200, 0, 65, color=image.COLOR_GREEN))
    buttons.append(create_button("Blue", 300, 0, 65, color=image.COLOR_BLUE))
    buttons.append(create_button("Show threshold", 400, 0))
    
    buttons.append(create_button("L+", 0, 440, 80, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("L-", 90, 440, 80, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("A+", 180, 440, 80, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("A-", 270, 440, 80, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("B+", 360, 440, 80, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("B-", 450, 440, 80, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("S", 540, 440, 40, 40, color=image.COLOR_YELLOW))
    buttons.append(create_button("L", 590, 440, 40, 40, color=image.COLOR_YELLOW))
    
    # load_thresholds()  # 加载阈值
    run()
